Retry with rbf = -1 if tokens are dropped - #5080
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a retry mechanism (retry_when_tokens_dropped) for Mixture of Experts (MoE) training, allowing the training step to be retried with a dropless buffer capacity if tokens are dropped due to ragged sort buffer overflow. The changes span configuration files, MoE routing layers, training loops, and unit tests. The review feedback highlights several critical issues: a performance regression caused by host-device synchronization on every step, a logging bug where metrics are overwritten during a retry, compatibility issues with pure NNX models due to the use of self.sow, a trivial unit test that does not actually execute the retry logic, and a JAX tracing anti-pattern when collecting intermediate metrics.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
6abb4a2 to
b074035
Compare
b074035 to
140bcdf
Compare
| # capacity by leveraging the helper _truncate_matrix. | ||
| local_group_size = _truncate_matrix(local_group_size[:, None], buffer_size)[:, 0] | ||
| if num_expert_parallelism > 1: | ||
| has_overflow = jax.lax.psum(local_overflow, self._expert_parallelism_name) > 0 |
There was a problem hiding this comment.
Don't we have to check all (FS)DP ranks as well? DP rank 0 might not overflow for any EP but DP rank 1 would overflow - these ranks will have different values for has_overflow. Can we just all reduce over every device?
| ) | ||
| self.assertEqual(donate_argnums, 0) | ||
|
|
||
| def test_donate_argnums_is_zero_with_retry_when_tokens_dropped(self): |
There was a problem hiding this comment.
Does flipping the value of retry_when_tokens_drop change the signature of train? I'm not sure whats being tested here that is specific to retry_when_tokens_dropped
140bcdf to
0672353
Compare
0672353 to
2173d5d
Compare
…rflow Implements host-driven step replay for handling token dropping with HBM OOM protection: - Pre-compiles both a candidate train step (using tuned ragged_buffer_factor) and a dropless train step (using dropless buffer). - Disables state donation (donate_argnums = ()) when retry_when_tokens_dropped is enabled so input state is preserved for replay. - In the training loop, evaluates metrics['has_moe_overflow'] on the host. If True, discards candidate state, deallocates its device buffers, and replays the step with the dropless executable. - All-reduces local_overflow across all mesh axes (tuple(self.mesh.axis_names)) inside shard_map to ensure all (FS)DP ranks observe identical overflow status, avoiding multi-host hangs. - Configures dropless fallback with token chunking, optimization barrier, and full remat to fit worst-case buffer allocations in HBM. - Eliminates jax.lax.cond control flow completely from the forward and backward graph, preserving compiler optimizations and cross-layer collective pipelining.
2173d5d to
2b2becc
Compare
2 participants
Description
Currently, when ragged buffer is set, we drop any tokens that exceed the buffer size. This can occur if the dataset is very imbalanced.
This PR implements an optional flag to prevent token dropping. If it's set to true, then we will rerun the given layer if tokens are dropped with worse-case ragged buffer size to prevent token dropping.
BUGS: b/555278394
Tests
We saw retry successfully retries and has similar convergence to baseline when ragged buffer drops tokens (and diverges when retry is disabled), and when tokens were not dropped (rbf = 2) throughput is similar to baseline. Baseline experiments were executed on
cl/974137029for direct comparison. All workloads evaluated on 256 TPU chips (DeepSeek-V3 671B MoE).p4head rbf = -12866460823.473(ppl: 32.24)631.26(mean step: 12.98s)deploy_p4head_7029_rbf_neg1p4head rbf = 2(No Retry)2866460893.481(ppl: 32.50)653.55(mean step: 12.53s)deploy_p4head_7029_rbf_2pr rbf = -1(No Retry)2867183163.457(ppl: 31.72)632.81(mean step: 12.95s)deploy_pr_6389_rbf_neg1pr rbf = 0.125(No Retry)2866460925.271(ppl: 194.57)949.74(mean step: 8.63s)deploy_cl_rbf_0p125_no_retrypr rbf = 0.125 with retry2866214023.477(train ppl: 32.35)391.64(mean step: 20.92s)deploy_cl_rbf_0p125_retrypr rbf = 2 with retry2867183403.468(ppl: 32.07)653.84(mean step: 12.53s)deploy_pr_6389_rbf_2_retrypr rbf = 2 without retry2867192683.482(ppl: 32.52)653.48(mean step: 12.54s)deploy_pr_6389_rbf_2_no_retryChecklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.